import os
import json
import glob
import argparse
classes = {'door':0,'window':1,'table':2,'hostage':3,'militant':4,'face':4,'maskface':6,'chair':7,'person':8}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--path',help='json path')
    args = parser.parse_args()
    return args

def json2yolo(path):
    jsonlist = glob.glob(path+'/*.json')
    for j in jsonlist:
        json_file = json.load(open(j, "r", encoding="utf-8"))
        imageHeight = json_file['imageHeight']
        imageWidth = json_file['imageWidth']
        f = open(j[:-4]+'txt','w')
        for target in json_file['shapes']:
            cls = target['label']
            x1,y1 = target['points'][0]
            x2,y2 = target['points'][1]
            cls = classes[cls]
            cx = (x1+x2)/2/imageWidth
            cy = (y1+y2)/2/imageHeight
            w = (x2-x1)/imageWidth
            h = (y2-y1)/imageHeight
            f.write(str(cls)+' '+str(cx)+' '+str(cy)+' '+str(w)+' '+str(h)+'\n')
        f.close()

if __name__ == '__main__':
    args = parse_args()
    json2yolo(args.path)
